import subprocess
import os

try:
    import huggingface_hub
except ImportError:
    print("huggingface_hub 库未安装，正在安装...")
    os.system("pip install huggingface_hub -i https://pypi.tuna.tsinghua.edu.cn/simple/")
    from huggingface_hub import list_repo_files


#from app import download_model_or_dataset, list_files
try:
    import gradio as gr
except ImportError:
    print("Gradio 库未安装，正在安装...")
    os.system("pip install gradio -i https://pypi.tuna.tsinghua.edu.cn/simple/")
    import gradio as gr

#显示模型/数据集清单
def list_files(repo_id, repo_type):
    try:
        return list_repo_files(repo_id=repo_id, repo_type=repo_type)
    except Exception as e:
        print(f"Error listing files for {repo_id}: {str(e)}")
        return []

#模型/数据集下载主程序
def download_model_or_dataset(model, token, include, exclude, dataset, save_dir, use_hf_transfer, use_mirror):
    token_option = f"--token {token}" if token else ""
    include_option = f"--include {include}" if include else ""
    exclude_option = f"--exclude {exclude}" if exclude else ""

    if model:
        model_name = model.split("/")
        save_dir_option = (
            f"--local-dir {os.path.join(save_dir, 'models--' + model_name[0] + '--' + model_name[1])}"
            if save_dir
            else ""
        )

        download_command = (
            f"huggingface-cli download {token_option} {include_option} {exclude_option} "
            f"--local-dir-use-symlinks False --resume-download {model} {save_dir_option}"
        )
    elif dataset:
        dataset_name = dataset.split("/")
        save_dir_option = (
            f"--local-dir {os.path.join(save_dir, 'datasets--' + dataset_name[0] + '--' + dataset_name[1])}"
            if save_dir
            else ""
        )

        download_command = (
            f"huggingface-cli download {token_option} {include_option} {exclude_option} "
            f"--local-dir-use-symlinks False --resume-download --repo-type dataset {dataset} {save_dir_option}"
        )

    if use_hf_transfer:
        os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
        print("HF transfer enabled")

    if use_mirror:
        os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
        print("Mirror enabled")

    subprocess.run(download_command)

#更新选择的模型/数据集清单
def update_file_list(repo_id, repo_type):
    files = list_files(repo_id, repo_type)
    return gr.update(choices=files, value=files)

#启动webui
def start_webui():
    with gr.Blocks() as interface:
        gr.Markdown("# Huggingface模型/数据集WebUI下载工具")

        model_tab = gr.Tab("模型设置")
        dataset_tab = gr.Tab("数据集设置")

        with model_tab:
            with gr.Column():
                model = gr.Textbox(label="模型名称", placeholder="模型 URL")
                更新按钮 = gr.Button("更新文件列表(需能直连Huggingface)")
                file_list = gr.CheckboxGroup(label="可下载文件", interactive=True)
        with dataset_tab:
            with gr.Column():
                dataset = gr.Textbox(label="数据集名称", placeholder="数据集 URL")

        with gr.Row():
            token = gr.Textbox(label="token", placeholder="HuggingFace的认证token(可选)")
        with gr.Row():
            include = gr.Textbox(label="包含的下载文件 (可选)", placeholder="要包含的文件列表，用逗号分隔")
        with gr.Row():
            exclude = gr.Textbox(label="排除的下载文件 (可选)", placeholder="要排除的文件列表，用逗号分隔")
        with gr.Row():
            save_dir = gr.Textbox(label="保存目录(可选)", placeholder="模型或数据集的保存路径")
        with gr.Row():
            use_hf_transfer = gr.Checkbox(label="使用hf_transfer加速下载", value=False)
            use_mirror = gr.Checkbox(label="从HuggingFace镜像站下载", value=True)

        with gr.Row():
            下载按钮 = gr.Button("下载")
            
        更新按钮.click(fn=update_file_list, inputs=[model, gr.Textbox(visible=False, value="model")], outputs=file_list)
        下载按钮.click(fn=download_model_or_dataset, inputs=[model, token, include, exclude, dataset, save_dir, use_hf_transfer, use_mirror], outputs=gr.HTML(value="<h1>下载开始</h1>"))

    interface.launch()

if __name__ == '__main__':
    start_webui()
